# Redis客户端模块 - 用于实时数据的短期存储

import redis
import json
import logging
from typing import Dict, List, Any, Optional, Union
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

class RedisClient:
    """
    Redis客户端类，用于实时数据的短期存储
    实现数据的2小时保留策略
    """
    
    def __init__(self, host: str = 'localhost', port: int = 6379, 
                 db: int = 1, password: str = None, ttl: int = 28800):
        """
        初始化Redis连接
        
        Args:
            host: Redis服务器地址
            port: Redis服务器端口
            db: 数据库编号
            password: 密码
            ttl: 数据过期时间（秒），默认2小时
        """
        self._connected = False
        try:
            self.redis = redis.Redis(
                host=host,
                port=port,
                db=db,
                password=password,
                decode_responses=True  # 自动将字节解码为字符串
            )
            self.ttl = ttl  # 默认2小时过期
            
            self.ping()  # 测试连接
            self._connected = True
            logger.info(f"Redis连接初始化成功: {host}:{port}/{db}")
        except Exception as e:
            logger.error(f"Redis连接初始化失败: {str(e)}")
            raise

    def ping(self) -> bool:
        """测试Redis连接"""
        try:
            self.redis.ping()
            self._connected = True
            return True
        except:
            self._connected = False
            return False
    
    def is_connected(self) -> bool:
        """获取连接状态"""
        if not self._connected:
            return self.ping()
        return self._connected

    def get(self, key: str) -> Optional[str]:
        """获取Redis数据"""
        try:
            data = self.redis.get(key)
            if data and key.startswith('stock_list:') and key != 'stock_list:update_time':
                return json.loads(data)  # 只对股票列表数据自动解析JSON，更新时间除外
            return data
        except Exception as e:
            logger.error(f"获取Redis数据失败: {str(e)}")
            return None

    def set(self, key: str, value: str):
        """设置Redis数据"""
        try:
            if key.startswith('stock_list:') and not isinstance(value, str):
                value = json.dumps(value)  # 对于股票列表数据自动转换为JSON
            self.redis.set(key, value)
        except Exception as e:
            logger.error(f"设置Redis数据失败: {str(e)}")
    
    def store_realtime_quote(self, symbol: str, quote_data: Dict):
        """
        存储实时行情数据
        
        Args:
            symbol: 股票代码
            quote_data: 行情数据字典
        """
        try:
            # 构建键名
            key = f"quote:{symbol}"
            
            # 添加时间戳
            quote_data['timestamp'] = datetime.now().strftime('%Y%m%d %H:%M:%S')
            
            # 存储数据
            self.set(key, json.dumps(quote_data))
            logger.debug(f"成功存储股票{symbol}的实时行情数据")
        except Exception as e:
            logger.error(f"存储实时行情数据失败: {str(e)}")
    
    def store_realtime_tick(self, symbol: str, tick_data: Dict):
        """
        存储实时逐笔交易数据
        
        Args:
            symbol: 股票代码
            tick_data: 逐笔交易数据字典
        """
        try:
            # 构建键名
            timestamp = datetime.now().strftime('%Y%m%d%H%M%S%f')[:-3]  # 精确到毫秒
            key = f"tick:{symbol}:{timestamp}"
            
            # 添加时间戳
            tick_data['timestamp'] = timestamp
            
            # 存储数据
            self.set(key, json.dumps(tick_data))
            
            # 将键添加到有序集合中，用于按时间顺序检索
            score = float(timestamp)
            self.redis.zadd(f"ticks:{symbol}", {key: score})
            self.redis.expire(f"ticks:{symbol}", self.ttl)
            
            logger.debug(f"成功存储股票{symbol}的实时逐笔交易数据")
        except Exception as e:
            logger.error(f"存储实时逐笔交易数据失败: {str(e)}")
    
    def store_realtime_order(self, symbol: str, order_data: Dict):
        """
        存储实时委托订单数据
        
        Args:
            symbol: 股票代码
            order_data: 委托订单数据字典
        """
        try:
            # 构建键名
            timestamp = datetime.now().strftime('%Y%m%d%H%M%S%f')[:-3]  # 精确到毫秒
            key = f"order:{symbol}:{timestamp}"
            
            # 添加时间戳
            order_data['timestamp'] = timestamp
            
            # 存储数据
            self.set(key, json.dumps(order_data))
            
            # 将键添加到有序集合中，用于按时间顺序检索
            score = float(timestamp)
            self.redis.zadd(f"orders:{symbol}", {key: score})
            self.redis.expire(f"orders:{symbol}", self.ttl)
            
            logger.debug(f"成功存储股票{symbol}的实时委托订单数据")
        except Exception as e:
            logger.error(f"存储实时委托订单数据失败: {str(e)}")
    
    def store_market_status(self, status_data: Dict):
        """
        存储市场状态数据
        
        Args:
            status_data: 市场状态数据字典
        """
        try:
            # 构建键名
            key = "market:status"
            
            # 添加时间戳
            status_data['timestamp'] = datetime.now().strftime('%Y%m%d %H:%M:%S')
            
            # 存储数据
            self.set(key, json.dumps(status_data))
            logger.debug("成功存储市场状态数据")
        except Exception as e:
            logger.error(f"存储市场状态数据失败: {str(e)}")
    
    def get_realtime_quote(self, symbol: str) -> Dict:
        """
        获取实时行情数据
        
        Args:
            symbol: 股票代码
            
        Returns:
            Dict: 行情数据字典
        """
        try:
            # 构建键名
            key = f"quote:{symbol}"
            
            # 获取数据
            data = self.get(key)
            if data:
                return json.loads(data)
            return {}
        except Exception as e:
            logger.error(f"获取实时行情数据失败: {str(e)}")
            return {}
    
    def get_realtime_ticks(self, symbol: str, count: int = 10) -> List[Dict]:
        """
        获取实时逐笔交易数据
        
        Args:
            symbol: 股票代码
            count: 返回的记录数量
            
        Returns:
            List[Dict]: 逐笔交易数据列表
        """
        try:
            # 构建键名
            key_pattern = f"ticks:{symbol}"
            
            # 获取最近的n条记录
            tick_keys = self.redis.zrevrange(key_pattern, 0, count - 1)
            
            # 获取数据
            result = []
            for tick_key in tick_keys:
                data = self.get(tick_key)
                if data:
                    result.append(json.loads(data))
            
            return result
        except Exception as e:
            logger.error(f"获取实时逐笔交易数据失败: {str(e)}")
            return []
    
    def get_market_status(self) -> Dict:
        """
        获取市场状态数据
        
        Returns:
            Dict: 市场状态数据字典
        """
        try:
            # 构建键名
            key = "market:status"
            
            # 获取数据
            data = self.get(key)
            if data:
                return json.loads(data)
            return {}
        except Exception as e:
            logger.error(f"获取市场状态数据失败: {str(e)}")
            return {}

    def publish_market_event(self, event_type: str, event_data: Dict):
        """
        发布市场事件
        
        Args:
            event_type: 事件类型
            event_data: 事件数据
        """
        try:
            # 构建消息
            message = {
                'type': event_type,
                'data': event_data,
                'timestamp': datetime.now().strftime('%Y%m%d %H:%M:%S')
            }
            
            # 发布消息
            self.redis.publish('market:events', json.dumps(message))
            logger.debug(f"成功发布市场事件: {event_type}")
        except Exception as e:
            logger.error(f"发布市场事件失败: {str(e)}")
    
    def clear_expired_data(self):
        """
        清理过期数据（超过TTL的数据）
        """
        try:
            # 获取所有键
            keys = self.redis.keys("quote:*") + self.redis.keys("tick:*") + \
                   self.redis.keys("order:*") + self.redis.keys("ticks:*") + \
                   self.redis.keys("orders:*")
            
            # 检查每个键的过期时间
            now = datetime.now()
            for key in keys:
                ttl = self.redis.ttl(key)
                if ttl < 0:  # 已过期或没有设置过期时间
                    self.redis.delete(key)
            
            logger.info("成功清理过期数据")
        except Exception as e:
            logger.error(f"清理过期数据失败: {str(e)}")
    
    def close(self):
        """
        关闭Redis连接
        """
        if hasattr(self, 'redis'):
            self.redis.close()
            logger.info("Redis连接已关闭")
    
    def __del__(self):
        """
        析构函数，确保连接被关闭
        """
        self.close()